from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from core.prompt_session import PromptSession
from dotenv import load_dotenv
from typing import List, Optional
import os

load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")

app = FastAPI()

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, replace with specific origins
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Store sessions in memory
sessions = {}

class PromptRequest(BaseModel):
    session_id: str
    user_input: str
    task: str = "general_assistant"
    model: str = "gpt-3.5-turbo"

class MemoryItem(BaseModel):
    content: str

class SessionMemory(BaseModel):
    session_id: str
    memory: List[str]
    task: str

@app.post("/chat")
def chat(request: PromptRequest):
    # Get or create session
    if request.session_id not in sessions:
        sessions[request.session_id] = PromptSession(
            task=request.task, 
            session_id=request.session_id
        )

    session = sessions[request.session_id]
    
    # Check if this is a memory command
    memory_response = session.add_user_input(request.user_input)
    
    # If it's a memory command, return the response directly
    if memory_response:
        return {"reply": memory_response, "is_memory_command": True}
    
    # Otherwise, get normal AI response
    try:
        reply = session.get_response(model=request.model)
        return {"reply": reply, "is_memory_command": False}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/memory/{session_id}")
def get_memory(session_id: str):
    """Get the memory for a specific session."""
    if session_id not in sessions:
        # Try to load from saved sessions
        try:
            temp_session = PromptSession(session_id=session_id)
            return {"session_id": session_id, "memory": temp_session.memory, "task": temp_session.task}
        except:
            raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
    
    session = sessions[session_id]
    return {"session_id": session_id, "memory": session.memory, "task": session.task}

@app.post("/memory/{session_id}")
def add_memory(session_id: str, memory_item: MemoryItem):
    """Add a memory item to a session."""
    if session_id not in sessions:
        # Try to load from saved sessions
        try:
            sessions[session_id] = PromptSession(session_id=session_id)
        except:
            raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
    
    session = sessions[session_id]
    session._add_to_memory(memory_item.content)
    session._save_memory()
    return {"status": "success", "message": "Memory added"}

@app.delete("/memory/{session_id}")
def clear_memory(session_id: str):
    """Clear all memory for a session."""
    if session_id not in sessions:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
    
    session = sessions[session_id]
    session.memory = []
    session._save_memory()
    return {"status": "success", "message": "Memory cleared"}

@app.delete("/memory/{session_id}/{index}")
def delete_memory_item(session_id: str, index: int):
    """Delete a specific memory item from a session."""
    if session_id not in sessions:
        raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
    
    session = sessions[session_id]
    if index < 0 or index >= len(session.memory):
        raise HTTPException(status_code=400, detail="Memory index out of range")
    
    deleted = session.memory.pop(index)
    session._save_memory()
    return {"status": "success", "message": f"Memory deleted: {deleted}"}

@app.get("/sessions")
def list_sessions():
    """List all available sessions."""
    session_info = []
    
    # List sessions from memory
    for session_id, session in sessions.items():
        session_info.append({
            "session_id": session_id,
            "task": session.task,
            "memory_count": len(session.memory)
        })
    
    # Also list sessions from disk
    import os
    memory_dir = "session_memories"
    if os.path.exists(memory_dir):
        for filename in os.listdir(memory_dir):
            if filename.endswith('.json'):
                session_id = filename[:-5]  # Remove .json extension
                if session_id not in [s["session_id"] for s in session_info]:
                    try:
                        temp_session = PromptSession(session_id=session_id)
                        session_info.append({
                            "session_id": session_id,
                            "task": temp_session.task,
                            "memory_count": len(temp_session.memory)
                        })
                    except:
                        pass
    
    return {"sessions": session_info}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)